{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6-3 使用单GPU训练模型\n",
    "\n",
    "深度学习的训练过程常常非常耗时，一个模型训练几个小时是家常便饭，训练几天也是常有的事情，有时候甚至要训练几十天。\n",
    "\n",
    "训练过程的耗时主要来自于两个部分，一部分来自数据准备，另一部分来自参数迭代。\n",
    "\n",
    "当数据准备过程还是模型训练时间的主要瓶颈时，我们可以使用更多进程来准备数据。\n",
    "\n",
    "当参数迭代过程成为训练时间的主要瓶颈时，我们通常的方法是应用GPU或者Google的TPU来进行加速。\n",
    "\n",
    "详见《用GPU加速Keras模型——Colab免费GPU使用攻略》\n",
    "\n",
    "https://zhuanlan.zhihu.com/p/68509398"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "无论是内置fit方法，还是自定义训练循环，从CPU切换成单GPU训练模型都是非常方便的，无需更改任何代码。当存在可用的GPU时，如果不特意指定device，tensorflow会自动优先选择使用GPU来创建张量和执行张量计算。\n",
    "\n",
    "但如果是在公司或者学校实验室的服务器环境，存在多个GPU和多个使用者时，为了不让单个同学的任务占用全部GPU资源导致其他同学无法使用（tensorflow默认获取全部GPU的全部内存资源权限，但实际上只使用一个GPU的部分资源），我们通常会在开头增加以下几行代码以控制每个任务使用的GPU编号和显存大小，以便其他同学也能够同时训练模型。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在Colab笔记本中：修改->笔记本设置->硬件加速器 中选择 GPU\n",
    "\n",
    "注：以下代码只能在Colab 上才能正确执行。\n",
    "\n",
    "可通过以下colab链接测试效果《tf_单GPU》：\n",
    "\n",
    "https://colab.research.google.com/drive/1r5dLoeJq5z01sU72BX2M5UiNSkuxsEFe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.1.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras import * \n",
    "\n",
    "#打印时间分割线\n",
    "@tf.function\n",
    "def printbar():\n",
    "    ts = tf.timestamp()\n",
    "    today_ts = ts%(24*60*60)\n",
    "\n",
    "    hour = tf.cast(today_ts//3600+8,tf.int32)%tf.constant(24)\n",
    "    minite = tf.cast((today_ts%3600)//60,tf.int32)\n",
    "    second = tf.cast(tf.floor(today_ts%60),tf.int32)\n",
    "    \n",
    "    def timeformat(m):\n",
    "        if tf.strings.length(tf.strings.format(\"{}\",m))==1:\n",
    "            return(tf.strings.format(\"0{}\",m))\n",
    "        else:\n",
    "            return(tf.strings.format(\"{}\",m))\n",
    "    \n",
    "    timestring = tf.strings.join([timeformat(hour),timeformat(minite),\n",
    "                timeformat(second)],separator = \":\")\n",
    "    tf.print(\"==========\"*8,end = \"\")\n",
    "    tf.print(timestring)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 一、 GPU设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpus = tf.config.list_physical_devices(\"GPU\")\n",
    "\n",
    "if gpus:\n",
    "    gpu0 = gpus[0] #如果有多个GPU，仅使用第0个GPU\n",
    "    tf.config.experimental.set_memory_growth(gpu0, True) #设置GPU显存用量按需使用\n",
    "    # 或者也可以设置GPU显存为固定使用量(例如：4G)\n",
    "    #tf.config.experimental.set_virtual_device_configuration(gpu0,\n",
    "    #    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)]) \n",
    "    tf.config.set_visible_devices([gpu0],\"GPU\") "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "比较GPU和CPU的计算速度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "printbar()\n",
    "with tf.device(\"/gpu:0\"):\n",
    "    tf.random.set_seed(0)\n",
    "    a = tf.random.uniform((10000,100),minval = 0,maxval = 3.0)\n",
    "    b = tf.random.uniform((100,100000),minval = 0,maxval = 3.0)\n",
    "    c = a@b\n",
    "    tf.print(tf.reduce_sum(tf.reduce_sum(c,axis = 0),axis=0))\n",
    "printbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================20:31:02\n",
      "2.24953795e+11\n",
      "================================================================================20:31:07\n"
     ]
    }
   ],
   "source": [
    "printbar()\n",
    "with tf.device(\"/cpu:0\"):\n",
    "    tf.random.set_seed(0)\n",
    "    a = tf.random.uniform((10000,100),minval = 0,maxval = 3.0)\n",
    "    b = tf.random.uniform((100,100000),minval = 0,maxval = 3.0)\n",
    "    c = a@b\n",
    "    tf.print(tf.reduce_sum(tf.reduce_sum(c,axis = 0),axis=0))\n",
    "printbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 二、准备数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_LEN = 300\n",
    "BATCH_SIZE = 32\n",
    "(x_train,y_train),(x_test,y_test) = datasets.reuters.load_data()\n",
    "x_train = preprocessing.sequence.pad_sequences(x_train,maxlen=MAX_LEN)\n",
    "x_test = preprocessing.sequence.pad_sequences(x_test,maxlen=MAX_LEN)\n",
    "\n",
    "MAX_WORDS = x_train.max()+1\n",
    "CAT_NUM = y_train.max()+1\n",
    "\n",
    "ds_train = tf.data.Dataset.from_tensor_slices((x_train,y_train)) \\\n",
    "          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \\\n",
    "          .prefetch(tf.data.experimental.AUTOTUNE).cache()\n",
    "   \n",
    "ds_test = tf.data.Dataset.from_tensor_slices((x_test,y_test)) \\\n",
    "          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \\\n",
    "          .prefetch(tf.data.experimental.AUTOTUNE).cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 三、定义模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding (Embedding)        (None, 300, 7)            216874    \n",
      "_________________________________________________________________\n",
      "conv1d (Conv1D)              (None, 296, 64)           2304      \n",
      "_________________________________________________________________\n",
      "max_pooling1d (MaxPooling1D) (None, 148, 64)           0         \n",
      "_________________________________________________________________\n",
      "conv1d_1 (Conv1D)            (None, 146, 32)           6176      \n",
      "_________________________________________________________________\n",
      "max_pooling1d_1 (MaxPooling1 (None, 73, 32)            0         \n",
      "_________________________________________________________________\n",
      "flatten (Flatten)            (None, 2336)              0         \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 46)                107502    \n",
      "=================================================================\n",
      "Total params: 332,856\n",
      "Trainable params: 332,856\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "tf.keras.backend.clear_session()\n",
    "\n",
    "def create_model():\n",
    "    \n",
    "    model = models.Sequential()\n",
    "\n",
    "    model.add(layers.Embedding(MAX_WORDS,7,input_length=MAX_LEN))\n",
    "    model.add(layers.Conv1D(filters = 64,kernel_size = 5,activation = \"relu\"))\n",
    "    model.add(layers.MaxPool1D(2))\n",
    "    model.add(layers.Conv1D(filters = 32,kernel_size = 3,activation = \"relu\"))\n",
    "    model.add(layers.MaxPool1D(2))\n",
    "    model.add(layers.Flatten())\n",
    "    model.add(layers.Dense(CAT_NUM,activation = \"softmax\"))\n",
    "    return(model)\n",
    "\n",
    "model = create_model()\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 四、训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================20:34:54\n",
      "Epoch=1,Loss:1.99821234,Accuracy:0.477955908,Valid Loss:1.65013993,Valid Accuracy:0.574799657\n",
      "\n",
      "================================================================================20:34:58\n",
      "Epoch=2,Loss:1.44956291,Accuracy:0.624471188,Valid Loss:1.49682689,Valid Accuracy:0.631344616\n",
      "\n",
      "================================================================================20:35:02\n",
      "Epoch=3,Loss:1.13833857,Accuracy:0.7069695,Valid Loss:1.49673975,Valid Accuracy:0.650044501\n",
      "\n",
      "================================================================================20:35:06\n",
      "Epoch=4,Loss:0.853759468,Accuracy:0.779336452,Valid Loss:1.66020906,Valid Accuracy:0.655387342\n",
      "\n",
      "================================================================================20:35:10\n",
      "Epoch=5,Loss:0.620348155,Accuracy:0.844689369,Valid Loss:1.92321694,Valid Accuracy:0.649154067\n",
      "\n",
      "================================================================================20:35:14\n",
      "Epoch=6,Loss:0.466241539,Accuracy:0.888666213,Valid Loss:2.16813254,Valid Accuracy:0.652715921\n",
      "\n",
      "================================================================================20:35:18\n",
      "Epoch=7,Loss:0.373969585,Accuracy:0.912602961,Valid Loss:2.29692721,Valid Accuracy:0.650935\n",
      "\n",
      "================================================================================20:35:22\n",
      "Epoch=8,Loss:0.316198081,Accuracy:0.924070358,Valid Loss:2.40828896,Valid Accuracy:0.647818327\n",
      "\n",
      "================================================================================20:35:26\n",
      "Epoch=9,Loss:0.279562235,Accuracy:0.93219775,Valid Loss:2.49524879,Valid Accuracy:0.651380241\n",
      "\n",
      "================================================================================20:35:31\n",
      "Epoch=10,Loss:0.252419889,Accuracy:0.936205745,Valid Loss:2.58140373,Valid Accuracy:0.646482646\n",
      "\n"
     ]
    }
   ],
   "source": [
    "optimizer = optimizers.Nadam()\n",
    "loss_func = losses.SparseCategoricalCrossentropy()\n",
    "\n",
    "train_loss = metrics.Mean(name='train_loss')\n",
    "train_metric = metrics.SparseCategoricalAccuracy(name='train_accuracy')\n",
    "\n",
    "valid_loss = metrics.Mean(name='valid_loss')\n",
    "valid_metric = metrics.SparseCategoricalAccuracy(name='valid_accuracy')\n",
    "\n",
    "@tf.function\n",
    "def train_step(model, features, labels):\n",
    "    with tf.GradientTape() as tape:\n",
    "        predictions = model(features,training = True)\n",
    "        loss = loss_func(labels, predictions)\n",
    "    gradients = tape.gradient(loss, model.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "\n",
    "    train_loss.update_state(loss)\n",
    "    train_metric.update_state(labels, predictions)\n",
    "    \n",
    "@tf.function\n",
    "def valid_step(model, features, labels):\n",
    "    predictions = model(features)\n",
    "    batch_loss = loss_func(labels, predictions)\n",
    "    valid_loss.update_state(batch_loss)\n",
    "    valid_metric.update_state(labels, predictions)\n",
    "    \n",
    "\n",
    "def train_model(model,ds_train,ds_valid,epochs):\n",
    "    for epoch in tf.range(1,epochs+1):\n",
    "        \n",
    "        for features, labels in ds_train:\n",
    "            train_step(model,features,labels)\n",
    "\n",
    "        for features, labels in ds_valid:\n",
    "            valid_step(model,features,labels)\n",
    "\n",
    "        logs = 'Epoch={},Loss:{},Accuracy:{},Valid Loss:{},Valid Accuracy:{}'\n",
    "        \n",
    "        if epoch%1 ==0:\n",
    "            printbar()\n",
    "            tf.print(tf.strings.format(logs,\n",
    "            (epoch,train_loss.result(),train_metric.result(),valid_loss.result(),valid_metric.result())))\n",
    "            tf.print(\"\")\n",
    "            \n",
    "        train_loss.reset_states()\n",
    "        valid_loss.reset_states()\n",
    "        train_metric.reset_states()\n",
    "        valid_metric.reset_states()\n",
    "\n",
    "train_model(model,ds_train,ds_test,10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
